
/**
  ******************************************************************************
  * Copyright 2021 The Microbee Authors. All Rights Reserved.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  * 
  * http://www.apache.org/licenses/LICENSE-2.0
  * 
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * 
  * @file       gp_control.c
  * @author     baiyang
  * @date       2021-8-29
  ******************************************************************************
  */

/*----------------------------------include-----------------------------------*/
#include "gp_control.h"

#include <common/gp_math/gp_mathlib.h>
#include <internal_error/internal_error.h>
/*-----------------------------------macro------------------------------------*/
// control default definitions
#define CORNER_ACCELERATION_RATIO   1.0/math_sqrtf(2.0)   // acceleration reduction to enable zero overshoot corners
/*----------------------------------typedef-----------------------------------*/

/*---------------------------------prototype----------------------------------*/

/*----------------------------------variable----------------------------------*/

/*-------------------------------------os-------------------------------------*/

/*----------------------------------function----------------------------------*/
// proportional controller with piecewise sqrt sections to constrain second derivative
float control_sqrt_controller(float error, float p, float second_ord_lim, float dt)
{
    float correction_rate;
    if (math_flt_negative(second_ord_lim) || math_flt_zero(second_ord_lim)) {
        // second order limit is zero or negative.
        correction_rate = error * p;
    } else if (math_flt_zero(p)) {
        // P term is zero but we have a second order limit.
        if (math_flt_positive(error)) {
            correction_rate = math_sqrtf(2.0 * second_ord_lim * (error));
        } else if (math_flt_negative(error)) {
            correction_rate = -math_sqrtf(2.0 * second_ord_lim * (-error));
        } else {
            correction_rate = 0.0;
        }
    } else {
        // Both the P and second order limit have been defined.
        const float linear_dist = second_ord_lim / sq(p);
        if (error > linear_dist) {
            correction_rate = math_sqrtf(2.0 * second_ord_lim * (error - (linear_dist / 2.0)));
        } else if (error < -linear_dist) {
            correction_rate = -math_sqrtf(2.0 * second_ord_lim * (-error - (linear_dist / 2.0)));
        } else {
            correction_rate = error * p;
        }
    }
    if (!math_flt_positive(dt)) {
        // this ensures we do not get small oscillations by over shooting the error correction in the last time step.
        return math_constrain_float(correction_rate, -fabsf(error) / dt, fabsf(error) / dt);
    } else {
        return correction_rate;
    }
}

// proportional controller with piecewise sqrt sections to constrain second derivative
Vector2f_t control_sqrt_controller2(const Vector2f_t* error, float p, float second_ord_lim, float dt)
{
    Vector2f_t output = {0};
    const float error_length = vec2_length(error);
    if (!math_flt_positive(error_length)) {
        return output;
    }

    const float correction_length = control_sqrt_controller(error_length, p, second_ord_lim, dt);

    vec2_mult(&output, error, correction_length / error_length);

    return output;
}

// inverse of the sqrt controller.  calculates the input (aka error) to the sqrt_controller required to achieve a given output
static float inv_sqrt_controller(float output, float p, float D_max)
{
    if (math_flt_positive(D_max) && math_flt_zero(p)) {
        return (output * output) / (2.0 * D_max);
    }
    if ((math_flt_negative(D_max) || math_flt_zero(D_max)) && !math_flt_zero(p)) {
        return output / p;
    }
    if ((math_flt_negative(D_max) || math_flt_zero(D_max)) && math_flt_zero(p)) {
        return 0.0;
    }

    // calculate the velocity at which we switch from calculating the stopping point using a linear function to a sqrt function
    const float linear_velocity = D_max / p;

    if (fabsf(output) < linear_velocity) {
        // if our current velocity is below the cross-over point we use a linear function
        return output / p;
    }

    const float linear_dist = D_max / sq(p);
    const float stopping_dist = (linear_dist * 0.5f) + sq(output) / (2.0 * D_max);
    return math_flt_positive(output) ? stopping_dist : -stopping_dist;
}

// update_vel_accel - single axis projection of velocity, vel, forwards in time based on a time step of dt and acceleration of accel.
// the velocity is not moved in the direction of limit if limit is not set to zero
// limit - specifies if the system is unable to continue to accelerate.
// vel_error - specifies the direction of the velocity error used in limit handling.
void control_update_vel_accel(float* vel, float accel, float dt, float limit, float vel_error)
{
    float delta_vel = accel * dt;

    // do not add delta_vel if it will increase the velocity error in the direction of limit
    // unless adding delta_vel will reduce vel towards zero
    if (math_flt_positive(delta_vel * limit) && math_flt_positive(vel_error * limit)) {
        if (math_flt_negative((*vel) * limit)) {
            delta_vel = math_constrain_float(delta_vel, -fabsf(*vel), fabsf(*vel));
        } else {
            delta_vel = 0.0f;
        }
    }
    *vel += delta_vel;
}

// update_pos_vel_accel - single axis projection of position and velocity, pos and vel, forwards in time based on a time step of dt and acceleration of accel.
// the position and velocity is not moved in the direction of limit if limit is not set to zero
// limit - specifies if the system is unable to continue to accelerate.
// pos_error and vel_error - specifies the direction of the velocity error used in limit handling.
void control_update_pos_vel_accel(postype_t* pos, float* vel, float accel, float dt, float limit, float pos_error, float vel_error)
{
    // move position and velocity forward by dt if it does not increase error when limited.
    float delta_pos = *vel * dt + accel * 0.5f * sq(dt);

    // do not add delta_pos if it will increase the velocity error in the direction of limit
    if (math_flt_positive(delta_pos * limit) && math_flt_positive(pos_error * limit)) {
        delta_pos = 0.0f;
    }
    *pos += delta_pos;

    control_update_vel_accel(vel, accel, dt, limit, vel_error);
}

// update_vel_accel - dual axis projection of position and velocity, pos and vel, forwards in time based on a time step of dt and acceleration of accel.
// the velocity is not moved in the direction of limit if limit is not set to zero.
// limit - specifies if the system is unable to continue to accelerate.
// pos_error and vel_error - specifies the direction of the velocity error used in limit handling.
void control_update_vel_accel_xy(Vector2f_t* vel, const Vector2f_t* accel, float dt, const Vector2f_t* limit, const Vector2f_t* vel_error)
{
    // increase velocity by acceleration * dt if it does not increase error when limited.
    // unless adding delta_vel will reduce the magnitude of vel
    Vector2f_t delta_vel = {accel->x * dt, accel->y * dt};
    if (!vec2_is_zero(limit) && !vec2_is_zero(&delta_vel)) {
        // check if delta_vel will increase the velocity error in the direction of limit
        if (math_flt_positive(vec2_dot(&delta_vel, limit)) && math_flt_positive(vec2_dot(vel_error, limit)) && !math_flt_negative(vec2_dot(vel, limit))) {
            vec2_zero(&delta_vel);
        }
    }

    vel->x += delta_vel.x;
    vel->y += delta_vel.y;
}

// update_pos_vel_accel - dual axis projection of position and velocity, pos and vel, forwards in time based on a time step of dt and acceleration of accel.
// the position and velocity is not moved in the direction of limit if limit is not set to zero
void control_update_pos_vel_accel_xy(Vector2p* pos, Vector2f_t* vel, const Vector2f_t* accel, float dt, const Vector2f_t* limit, const Vector2f_t* pos_error, const Vector2f_t* vel_error)
{
    // move position and velocity forward by dt.
    Vector2f_t delta_pos = {vel->x * dt + accel->x * 0.5f * sq(dt),
                            vel->y * dt + accel->y * 0.5f * sq(dt)};

    if (!math_flt_zero(vec2_length_squared(limit))) {
        // zero delta_pos if it will increase the velocity error in the direction of limit
        if (math_flt_positive(vec2_dot(&delta_pos, limit)) && math_flt_positive(vec2_dot(pos_error, limit))) {
            vec2_zero(&delta_pos);
        }
    }

    pos->x += delta_pos.x;
    pos->y += delta_pos.y;

    control_update_vel_accel_xy(vel, accel, dt, limit, vel_error);
}

/* shape_accel calculates a jerk limited path from the current acceleration to an input acceleration.
 The function takes the current acceleration and calculates the required jerk limited adjustment to the acceleration for the next time dt.
 The kinematic path is constrained by :
    maximum jerk - jerk_max (must be positive).
 The function alters the variable accel to follow a jerk limited kinematic path to accel_input.
*/
void control_shape_accel(float accel_input, float* accel,
                 float jerk_max, float dt)
{
    if (!math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    if (math_flt_positive(dt)) {
        // jerk limit acceleration change
        float accel_delta = accel_input - *accel;
        accel_delta = math_constrain_float(accel_delta, -jerk_max * dt, jerk_max * dt);
        *accel += accel_delta;
    }
}

// 2D version
void control_shape_accel_xy(const Vector2f_t* accel_input, Vector2f_t* accel,
                    float jerk_max, float dt)
{
    if (!math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    if (math_flt_positive(dt)) {
        // jerk limit acceleration change
        Vector2f_t accel_delta = {accel_input->x - accel->x, accel_input->y - accel->y};

        vec2_limit_length(&accel_delta, jerk_max * dt);

        accel->x = accel->x + accel_delta.x;
        accel->y = accel->y + accel_delta.y;
    }
}

void control_shape_accel_xy2(const Vector3f_t* accel_input, Vector3f_t* accel,
                    float jerk_max, float dt)
{
    const Vector2f_t accel_input_2f  = {accel_input->x, accel_input->y};
    Vector2f_t accel_2f  = {accel->x, accel->y};

    control_shape_accel_xy(&accel_input_2f, &accel_2f, jerk_max, dt);
    accel->x = accel_2f.x;
    accel->y = accel_2f.y;
}


/* shape_vel_accel and shape_vel_xy calculate a jerk limited path from the current position, velocity and acceleration to an input velocity.
 The function takes the current position, velocity, and acceleration and calculates the required jerk limited adjustment to the acceleration for the next time dt.
 The kinematic path is constrained by :
    minimum acceleration - accel_min (must be negative),
    maximum acceleration - accel_max (must be positive),
    maximum jerk - jerk_max (must be positive).
 The function alters the variable accel to follow a jerk limited kinematic path to vel_input and accel_input.
 The correction acceleration is limited from accel_min to accel_max. If limit_total is true the target acceleration is limited from accel_min to accel_max.
*/
void control_shape_vel_accel(float vel_input, float accel_input,
                     float vel, float* accel,
                     float accel_min, float accel_max,
                     float jerk_max, float dt, bool limit_total_accel)
{
    // sanity check accel_max
    if (!math_flt_negative(accel_min) || !math_flt_positive(accel_max) || !math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    // velocity error to be corrected
    float vel_error = vel_input - vel;

    // Calculate time constants and limits to ensure stable operation
    float KPa;

    // The direction of acceleration limit is the same as the velocity error.
    // This is because the velocity error is negative when slowing down while
    // closing a positive position error.
    if (math_flt_positive(vel_error)) {
        KPa = jerk_max / accel_max;
    } else {
        KPa = jerk_max / (-accel_min);
    }

    // acceleration to correct velocity
    float accel_target = control_sqrt_controller(vel_error, KPa, jerk_max, dt);

    // constrain correction acceleration from accel_min to accel_max
    accel_target = math_constrain_float(accel_target, accel_min, accel_max);

    // velocity correction with input velocity
    accel_target += accel_input;

    // constrain total acceleration from accel_min to accel_max
    if (limit_total_accel) {
        accel_target = math_constrain_float(accel_target, accel_min, accel_max);
    }

    control_shape_accel(accel_target, accel, jerk_max, dt);
}

// 2D version
void control_shape_vel_accel_xy(const Vector2f_t *vel_input1, const Vector2f_t* accel_input,
                        const Vector2f_t* vel, Vector2f_t* accel,
                        float accel_max, float jerk_max, float dt, bool limit_total_accel)
{
    // sanity check accel_max
    if (!math_flt_positive(accel_max) || !math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    Vector2f_t vel_input = *vel_input1;

    // Calculate time constants and limits to ensure stable operation
    const float KPa = jerk_max / accel_max;

    // velocity error to be corrected
    const Vector2f_t vel_error = {vel_input.x - vel->x, vel_input.y - vel->y};

    // acceleration to correct velocity
    Vector2f_t accel_target = control_sqrt_controller2(&vel_error, KPa, jerk_max, dt);

    // limit correction acceleration to accel_max
    vec2_limit_length(&accel_target, accel_max);

    // limit correction acceleration to accel_max
    if (vec2_is_zero(&vel_input)) {
        vec2_limit_length(&accel_target, accel_max);
    } else {
        // calculate acceleration in the direction of and perpendicular to the velocity input
        const Vector2f_t vel_input_unit = vec2_normalized(&vel_input);
        float accel_dir = vec2_cross(&vel_input_unit, &accel_target);
        Vector2f_t accel_cross = {accel_target.x - (vel_input_unit.x * accel_dir),
                                  accel_target.y - (vel_input_unit.y * accel_dir)};

        // ensure 1/sqrt(2) of maximum acceleration is available to correct cross component 
        // relative to vel_input
        if (sq(accel_dir) <= vec2_length_squared(&accel_cross)) {
            // accel_target can be simply limited in magnitude
            vec2_limit_length(&accel_target, accel_max);
        } else {
            // limiting the length of the vector will reduce the lateral acceleration below 1/sqrt(2)
            // limit the lateral acceleration to 1/sqrt(2) and retain as much of the remaining
            // acceleration as possible.
            vec2_limit_length(&accel_cross, CORNER_ACCELERATION_RATIO * accel_max);
            float accel_max_dir = math_sqrtf(sq(accel_max) - vec2_length_squared(&accel_cross));
            accel_dir = math_constrain_float(accel_dir, -accel_max_dir, accel_max_dir);

            accel_target.x = accel_cross.x + vel_input_unit.x * accel_dir;
            accel_target.y = accel_cross.y + vel_input_unit.y * accel_dir;
        }
    }

    vec2_add(&accel_target, &accel_target, accel_input);

    // limit total acceleration to accel_max
    if (limit_total_accel) {
        vec2_limit_length(&accel_target, accel_max);
    }

    control_shape_accel_xy(&accel_target, accel, jerk_max, dt);
}

/* shape_pos_vel_accel calculate a jerk limited path from the current position, velocity and acceleration to an input position and velocity.
 The function takes the current position, velocity, and acceleration and calculates the required jerk limited adjustment to the acceleration for the next time dt.
 The kinematic path is constrained by :
    minimum velocity - vel_min (must not be positive),
    maximum velocity - vel_max (must not be negative),
    minimum acceleration - accel_min (must be negative),
    maximum acceleration - accel_max (must be positive),
    maximum jerk - jerk_max (must be positive).
 The function alters the variable accel to follow a jerk limited kinematic path to pos_input, vel_input and accel_input.
 The correction velocity is limited to vel_max to vel_min. If limit_total is true the target velocity is limited to vel_max to vel_min.
 The correction acceleration is limited from accel_min to accel_max. If limit_total is true the target acceleration is limited from accel_min to accel_max.
*/
void control_shape_pos_vel_accel(postype_t pos_input, float vel_input, float accel_input,
                         postype_t pos, float vel, float* accel,
                         float vel_min, float vel_max,
                         float accel_min, float accel_max,
                         float jerk_max, float dt, bool limit_total)
{
    // sanity check accel_max
    if (math_flt_positive(vel_min) || math_flt_negative(vel_max) || !math_flt_negative(accel_min) || !math_flt_positive(accel_max) || !math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    // position error to be corrected
    float pos_error = pos_input - pos;

    float accel_tc_max;
    float KPv;
    // Calculate time constants and limits to ensure stable operation
    // The negative acceleration limit is used here because the square root controller
    // manages the approach to the setpoint. Therefore the acceleration is in the opposite
    // direction to the position error.
    if (math_flt_positive(pos_error)) {
        accel_tc_max = -0.5 * accel_min;
        KPv = 0.5 * jerk_max / (-accel_min);
    } else {
        accel_tc_max = 0.5 * accel_max;
        KPv = 0.5 * jerk_max / accel_max;
    }

    // velocity to correct position
    float vel_target = control_sqrt_controller(pos_error, KPv, accel_tc_max, dt);

    // limit velocity to vel_max
    if (math_flt_negative(vel_min) || math_flt_positive(vel_max)) {
        vel_target = math_constrain_float(vel_target, vel_min, vel_max);
    }

    // velocity correction with input velocity
    vel_target += vel_input;

    // limit velocity between vel_min and vel_max
    if (limit_total) {
        vel_target = math_constrain_float(vel_target, vel_min, vel_max);
    }

    control_shape_vel_accel(vel_target, accel_input, vel, accel, accel_min, accel_max, jerk_max, dt, limit_total);
}

// 2D version
void control_shape_pos_vel_accel_xy(const Vector2p* pos_input, const Vector2f_t* vel_input, const Vector2f_t* accel_input,
                            const Vector2p* pos, const Vector2f_t* vel, Vector2f_t* accel,
                            float vel_max, float accel_max,
                            float jerk_max, float dt, bool limit_total)
{
    // sanity check accel_max
    if (math_flt_negative(vel_max) || !math_flt_positive(accel_max) || !math_flt_positive(jerk_max)) {
        INTERNAL_ERROR(internal_error_invalid_arg_or_result);
        return;
    }

    // Calculate time constants and limits to ensure stable operation
    const float KPv = 0.5 * jerk_max / accel_max;
    const float accel_tc_max = 0.5 * accel_max;

    // position error to be corrected
    Vector2f_t pos_error = {pos_input->x - pos->x, pos_input->y - pos->y};

    // velocity to correct position
    Vector2f_t vel_target = control_sqrt_controller2(&pos_error, KPv, accel_tc_max, dt);

    // limit velocity to vel_max
    if (math_flt_positive(vel_max)) {
        vec2_limit_length(&vel_target, vel_max);
    }

    // velocity correction with input velocity
    vec2_add(&vel_target, &vel_target, vel_input);

    // limit velocity to vel_max
    if (limit_total) {
        vec2_limit_length(&vel_target, vel_max);
    }

    control_shape_vel_accel_xy(&vel_target, accel_input, vel, accel, accel_max, jerk_max, dt, limit_total);
}

/* limit_accel_xy limits the acceleration to prioritise acceleration perpendicular to the provided velocity vector.
 Input parameters are:
    vel is the velocity vector used to define the direction acceleration limit is biased in.
    accel is the acceleration vector to be limited.
    accel_max is the maximum length of the acceleration vector after being limited.
 Returns true when accel vector has been limited.
*/
bool control_limit_accel_xy(const Vector2f_t* vel, Vector2f_t* accel, float accel_max)
{
    // check accel_max is defined
    if (!math_flt_positive(accel_max)) {
        return false;
    }
    // limit acceleration to accel_max while prioritizing cross track acceleration
    if (vec2_length_squared(accel) > sq(accel_max)) {
        if (vec2_is_zero(vel)) {
            // We do not have a direction of travel so do a simple vector length limit
            vec2_limit_length(accel, accel_max);
        } else {
            // calculate acceleration in the direction of and perpendicular to the velocity input
            const Vector2f_t vel_input_unit = vec2_normalized(vel);
            // acceleration in the direction of travel
            float accel_dir = vec2_dot(&vel_input_unit, accel);
            // cross track acceleration
            Vector2f_t accel_cross = {accel->x - (vel_input_unit.x * accel_dir),
                                      accel->y - (vel_input_unit.y * accel_dir)};
            if (vec2_limit_length(&accel_cross, accel_max)) {
                accel_dir = 0.0;
            } else {
                float accel_max_dir = math_sqrtf(sq(accel_max) - vec2_length_squared(&accel_cross));
                accel_dir = math_constrain_float(accel_dir, -accel_max_dir, accel_max_dir);
            }
            accel->x = accel_cross.x + vel_input_unit.x * accel_dir;
            accel->y = accel_cross.y + vel_input_unit.y * accel_dir;
        }
        return true;
    }
    return false;
}

// calculate the stopping distance for the square root controller based deceleration path
float control_stopping_distance(float velocity, float p, float accel_max)
{
    return inv_sqrt_controller(velocity, p, accel_max);
}

// calculate the maximum acceleration or velocity in a given direction
// based on horizontal and vertical limits.
float control_kinematic_limit(Vector3f_t direction, float max_xy, float max_z_pos, float max_z_neg)
{
    if (math_flt_zero(vec3_length_squared(&direction)) 
        || math_flt_zero(max_xy) 
        || math_flt_zero(max_z_pos) 
        || math_flt_zero(max_z_neg)) {
        return 0.0f;
    }

    max_xy = fabsf(max_xy);
    max_z_pos = fabsf(max_z_pos);
    max_z_neg = fabsf(max_z_neg);

    vec3_norm(&direction, &direction);

    Vector2f_t direction_tmp = {direction.x, direction.y};
    const float xy_length = vec2_length(&direction_tmp);

    if (math_flt_zero(xy_length)) {
        return math_flt_positive(direction.z) ? max_z_pos : max_z_neg;
    }

    if (math_flt_zero(direction.z)) {
        return max_xy;
    }

    const float slope = direction.z/xy_length;
    if (math_flt_positive(slope)) {
        if (fabsf(slope) < max_z_pos/max_xy) {
            return max_xy/xy_length;
        }
        return fabsf(max_z_pos/direction.z);
    }

    if (fabsf(slope) < max_z_neg/max_xy) {
        return max_xy/xy_length;
    }
    return fabsf(max_z_neg/direction.z);
}

// input_expo calculates the expo function on the normalised input.
// The input must be in the range of -1 to 1.
// The expo should be less than 1.0 but limited to be less than 0.95.
float control_input_expo(float input, float expo)
{
    input = math_constrain_float(input, -1.0, 1.0);
    if (expo < 0.95) {
        return (1 - expo) * input / (1 - expo * fabsf(input));
    }
    return input;
}

/*
  convert a maximum lean angle in degrees to an accel limit in m/s/s
 */
float control_angle_to_accel(float angle_deg)
{
    return GRAVITY_MSS * tanf(radians(angle_deg));
}

/*
  convert a maximum accel in m/s/s to a lean angle in degrees
 */
float control_accel_to_angle(float accel)
{
    return degrees(atanf((accel/GRAVITY_MSS)));
}

// control_rc_input_to_roll_pitch - transform pilot's normalised roll or pitch stick input into a roll and pitch euler angle command
// roll_in_unit and pitch_in_unit - are normalised roll and pitch stick input
// angle_max_deg - maximum lean angle from the z axis
// angle_limit_deg - provides the ability to reduce the maximum output lean angle to less than angle_max_deg
// returns roll and pitch angle in degrees
void control_rc_input_to_roll_pitch(float roll_in_unit, float pitch_in_unit, float angle_max_deg, float angle_limit_deg, float *roll_out_deg, float *pitch_out_deg)
{
    angle_max_deg = MIN(angle_max_deg, 85.0);
    float rc_2_rad = radians(angle_max_deg);

    // fetch roll and pitch stick positions and convert them to normalised horizontal thrust
    Vector2f_t thrust;
    thrust.x = - tanf(rc_2_rad * pitch_in_unit);
    thrust.y = tanf(rc_2_rad * roll_in_unit);

    // calculate the horizontal thrust limit based on the angle limit
    angle_limit_deg = math_constrain_float(angle_limit_deg, 10.0f, angle_max_deg);
    float thrust_limit = tanf(radians(angle_limit_deg));

    // apply horizontal thrust limit
    vec2_limit_length(&thrust, thrust_limit);

    // Conversion from angular thrust vector to euler angles.
    float pitch_rad = - atanf(thrust.x);
    float roll_rad = atanf(cosf(pitch_rad) * thrust.y);

    // Convert to degrees
    *roll_out_deg = degrees(roll_rad);
    *pitch_out_deg = degrees(pitch_rad);
}
/*------------------------------------test------------------------------------*/


